# test.py
import torch
import time
from ppo import PPO
from parallel_env import ParallelEnv
from environment import rm_simple_simulator
import os
import imageio
import numpy as np


if __name__ == "__main__":
    print("Initializing test environment...")
    # Create a dummy env for PPO initialization (n_envs=1)
    dummy_env = ParallelEnv(n_envs=1)
    ppo = PPO(dummy_env, lr=3e-4)  # lr doesn't matter for testing

    model_path = "model_checkpoint.pth"
    ppo.load(model_path)

    num_test_episodes = 3
    max_test_steps = 1000
    os.makedirs("test_gif", exist_ok=True)

    for test_idx in range(num_test_episodes):
        print(f"Running test episode {test_idx + 1}/{num_test_episodes}")
        test_env = rm_simple_simulator(render_mode="rgb_array")
        obs, _ = test_env.reset()
        frames = []
        done = False
        step_count = 0

        while not done and step_count < max_test_steps:
            actions, _, _ = ppo.policy.act(obs["visions"], obs["infos"])
            actions_np = (
                actions.cpu().detach().numpy().flatten()
            )  # (18,) for single env
            obs, rewards, dones, _, _ = test_env.step(actions_np)
            frame = test_env.render()
            if frame is not None:
                frames.append(frame)
            done = dones.any()
            step_count += 1
            time.sleep(0.2)  # To slow down rendering for visualization

        # Save GIF
        gif_path = f"test_gif/test_episode_{test_idx + 1}.gif"
        if frames:
            imageio.mimsave(gif_path, frames, fps=10)
            print(f"Saved test GIF: {gif_path} ({len(frames)} frames)")
        else:
            print(f"No frames captured for episode {test_idx + 1}")

        test_env.close()

    dummy_env.close()
    print("Testing completed.")
